package com.fy.fyai.common.utils;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

@Slf4j
@Component
public class OllamaEmbeddingUtil {
    private static final String OLLAMA_EMBEDDING_URL = "http://localhost:11434/api/embeddings";
    private static final String MODEL = "nomic-embed-text"; // 可根据实际模型名调整
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    public float[] getEmbedding(String text) {
        try {
            // 清理文本：移除换行符、多余空格，确保JSON格式正确
            String cleanText = text.replaceAll("\\s+", " ").trim();
            
            RestTemplate restTemplate = new RestTemplate();
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            String body = String.format("{\"model\":\"%s\",\"prompt\":\"%s\"}", MODEL, cleanText.replace("\"", "\\\""));
            HttpEntity<String> entity = new HttpEntity<>(body, headers);
            ResponseEntity<String> response = restTemplate.postForEntity(OLLAMA_EMBEDDING_URL, entity, String.class);
            if (response.getStatusCode().is2xxSuccessful()) {
                JsonNode root = OBJECT_MAPPER.readTree(response.getBody());
                JsonNode embeddingNode = root.get("embedding");
                if (embeddingNode != null && embeddingNode.isArray()) {
                    float[] vector = new float[embeddingNode.size()];
                    for (int i = 0; i < embeddingNode.size(); i++) {
                        vector[i] = (float) embeddingNode.get(i).asDouble();
                    }
                    return vector;
                }
            }
        } catch (Exception e) {
            log.error("获取embedding失败", e);
        }
        return new float[0];
    }
} 